"""
Doubao provider implementation.
"""

import os
import asyncio
import base64
import logging
from typing import Optional

from openai import OpenAI
from ..llm_provider import LLMProvider

# Set up logger
logger = logging.getLogger("droidrun")

class DoubaoProvider(LLMProvider):
    """Doubao provider implementation."""
    
    def _initialize_client(self) -> None:
        """Initialize the Doubao client."""
        # 从环境变量获取 API key、base_url、model_name
        self.api_key = self.api_key or os.environ.get("DOUBAO_API_KEY")
        self.base_url = self.base_url or os.environ.get("DOUBAO_BASE_URL")
        self.model_name = self.model_name or os.environ.get("DOUBAO_MODEL_NAME")
        if not self.api_key:
            raise ValueError("Doubao API key 未提供，也未在环境变量中找到")
        if not self.base_url:
            self.base_url = "https://openai.api.doubao.com/v1"  # 默认值，可根据实际情况修改
        if not self.model_name:
            self.model_name = "doubao-gpt-4"
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
        logger.info(f"Initialized Doubao client with model {self.model_name} at {self.base_url}")
    
    async def generate_response(
        self,
        system_prompt: str,
        user_prompt: str,
        screenshot_data: Optional[bytes] = None
    ) -> str:
        """使用 Doubao 生成回复。"""
        try:
            messages = [
                {"role": "system", "content": system_prompt},
            ]
            # 如果有截图，添加图片内容
            if screenshot_data:
                base64_image = base64.b64encode(screenshot_data).decode('utf-8')
                messages.append({
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Here's the current screenshot of the device. Please analyze it to help with the next action."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                })
            # 添加主用户 prompt
            messages.append({"role": "user", "content": user_prompt})
            response = await asyncio.to_thread(
                self.client.chat.completions.create,
                model=self.model_name,
                messages=messages,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                response_format={"type": "json_object"}
            )
            # 统计 token 用量
            usage = response.usage
            self.update_token_usage(usage.prompt_tokens, usage.completion_tokens)
            logger.info("===== Token Usage Statistics =====")
            logger.info(f"API Call #{self.api_calls}")
            logger.info(f"This call: {usage.prompt_tokens} prompt + {usage.completion_tokens} completion = {usage.total_tokens} tokens")
            logger.info(f"Cumulative: {self.get_token_usage_stats()}")
            logger.info("=================================")
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"Error calling Doubao API: {e}")
            raise 